import torch.nn as nn

from .embedding import BERTEmbedding
from .transformer import TransformerBlock


class BERT(nn.Module):
    """
    BERT 模型：来自 Transformers 的双向编码器表示。
    """

    def __init__(self, vocab_size, hidden=768, n_layers=12, attn_heads=12, dropout=0.1):
        """
        :param vocab_size: 总词汇量大小
        :param hidden: BERT 模型隐藏层大小
        :param n_layers: Transformer 块（层）的数量
        :param attn_heads: 注意力头的数量
        :param dropout: dropout 率
        """

        super().__init__()
        self.hidden = hidden
        self.n_layers = n_layers
        self.attn_heads = attn_heads

        # 论文指出他们使用 4*hidden_size 作为前馈网络隐藏层大小
        self.feed_forward_hidden = hidden * 4

        # BERT 的嵌入层，是位置嵌入、段嵌入和词嵌入的总和
        self.embedding = BERTEmbedding(vocab_size=vocab_size, embed_size=hidden)

        # 多层 Transformer 块，深度网络
        self.transformer_blocks = nn.ModuleList(
            [TransformerBlock(hidden, attn_heads, hidden * 4, dropout) for _ in range(n_layers)])

    def forward(self, x, segment_info):
        # 对 padding token 进行 attention masking
        # torch.ByteTensor([batch_size, 1, seq_len, seq_len)
        mask = (x > 0).unsqueeze(1).repeat(1, x.size(1), 1).unsqueeze(1)

        # 将索引序列嵌入为向量序列
        x = self.embedding(x, segment_info)

        # 在多个 Transformer 块上运行
        for transformer in self.transformer_blocks:
            x = transformer.forward(x, mask)

        return x
